// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/core/plan_node.h>
#include <pollux/exec/split.h>
#include <pollux/plan/reference_query_runner.h>
#include <pollux/testing/exec/util/query_assertions.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    const std::string kHiveConnectorId = "test-hive";

    struct SortingKeyAndOrder {
        const std::string key_;
        const core::SortOrder sortOrder_;

        SortingKeyAndOrder(std::string key, core::SortOrder sortOrder)
            : key_(std::move(key)), sortOrder_(std::move(sortOrder)) {
        }
    };

    /// Write the vector to the path.
    void writeToFile(
        const std::string &path,
        const VectorPtr &vector,
        memory::MemoryPool *pool);

    /// Write vectors into the path, one file per vector, and create spilts from
    /// those files.
    std::vector<Split> makeSplits(
        const std::vector<RowVectorPtr> &inputs,
        const std::string &path,
        const std::shared_ptr<memory::MemoryPool> &writerPool);

    /// Create splits with schema information from a directory.
    /// For example under directory /table1:
    /// /table1/p0=0/p1=0/00000_file1
    /// /table1/p0=0/p1=1/00000_file1
    /// It would return splits:
    /// split1 with partition keys (p0, 0), (p1, 0)
    /// split2 with partition keys (p0, 0), (p1, 1)
    std::vector<Split> makeSplits(const std::string &directory);

    /// Create a split from an exsiting file.
    Split makeSplit(
        const std::string &filePath,
        const std::unordered_map<std::string, std::optional<std::string> > &
                partitionKeys = {},
        std::optional<int32_t> tableBucketNumber = std::nullopt);

    /// Create a connector split from an exsiting file.
    std::shared_ptr<connector::ConnectorSplit> makeConnectorSplit(
        const std::string &filePath,
        const std::unordered_map<std::string, std::optional<std::string> > &
                partitionKeys = {},
        std::optional<int32_t> tableBucketNumber = std::nullopt);

    /// Create column names with the pattern '${prefix}${i}'.
    std::vector<std::string> makeNames(const std::string &prefix, size_t n);

    /// Create a batch consists of single all-null BIGINT column with as many rows
    /// as original input. Used when the query doesn't need to read any columns, but
    /// it needs to see a specific number of rows. This way we will be able to
    /// create a temporary test table with the necessary number of rows.
    RowVectorPtr makeNullRows(
        const std::vector<pollux::RowVectorPtr> &input,
        const std::string &colName,
        memory::MemoryPool *pool);

    /// Returns whether type is supported in TableScan. Empty Row type and Unknown
    /// type are not supported.
    bool isTableScanSupported(const TypePtr &type);

    /// Concat two RowTypes.
    RowTypePtr concat(const RowTypePtr &a, const RowTypePtr &b);

    /// Skip queries that use Timestamp, Varbinary, and IntervalDayTime types.
    /// DuckDB doesn't support nanosecond precision for timestamps or casting from
    /// Bigint to Interval.
    ///
    /// TODO Investigate mismatches reported when comparing Varbinary.
    bool containsUnsupportedTypes(const TypePtr &type);

    /// Determines whether the signature has an argument that contains typeName.
    /// typeName should be in lower case.
    bool usesInputTypeName(
        const exec::FunctionSignature &signature,
        const std::string &typeName);

    /// Determines whether the signature has an argument or return type that
    /// contains typeName. typeName should be in lower case.
    bool usesTypeName(
        const exec::FunctionSignature &signature,
        const std::string &typeName);

    // First resolves typeSignature. Then, the resolved type is a RowType or
    // contains RowTypes with empty field names, adds default names to these fields
    // in the RowTypes.
    TypePtr sanitizeTryResolveType(
        const exec::TypeSignature &typeSignature,
        const melon::F14FastMap<std::string, SignatureVariable> &variables,
        const melon::F14FastMap<std::string, TypePtr> &resolvedTypeVariables);

    TypePtr sanitizeTryResolveType(
        const exec::TypeSignature &typeSignature,
        const melon::F14FastMap<std::string, SignatureVariable> &variables,
        const melon::F14FastMap<std::string, TypePtr> &typeVariablesBindings,
        melon::F14FastMap<std::string, int> &integerVariablesBindings);

    // Invoked to set up memory system with arbitration.
    void setupMemory(
        int64_t allocatorCapacity,
        int64_t arbitratorCapacity,
        bool enableGlobalArbitration = true);

    /// Registers hive connector with configs. It should be called in the
    /// constructor of fuzzers that test plans with TableScan or uses
    /// PrestoQueryRunner that writes data to a local file.
    void registerHiveConnector(
        const std::unordered_map<std::string, std::string> &hiveConfigs);

    // Returns a PrestoQueryRunner instance if prestoUrl is non-empty. Otherwise,
    // returns a DuckQueryRunner instance and set disabled aggregation functions
    // properly.
    std::unique_ptr<ReferenceQueryRunner> setupReferenceQueryRunner(
        memory::MemoryPool *aggregatePool,
        const std::string &prestoUrl,
        const std::string &runnerName,
        const uint32_t &reqTimeoutMs);

    // Logs the input vectors if verbose logging is turned on.
    void logVectors(const std::vector<RowVectorPtr> &vectors);

    // Converts 'plan' into an SQL query and runs in the reference DB.
    // Result is returned as a MaterializedRowMultiset with the
    // ReferenceQueryErrorCode::kSuccess if successful, or an std::nullopt with a
    // ReferenceQueryErrorCode if the query fails.
    std::pair<std::optional<MaterializedRowMultiset>, ReferenceQueryErrorCode>
    computeReferenceResults(
        const core::PlanNodePtr &plan,
        ReferenceQueryRunner *referenceQueryRunner);

    // Similar to computeReferenceResults(), but returns the result as a
    // std::vector<RowVectorPtr>. This API throws if referenceQueryRunner doesn't
    // support returning results as a vector.
    std::pair<std::optional<std::vector<RowVectorPtr> >, ReferenceQueryErrorCode>
    computeReferenceResultsAsVector(
        const core::PlanNodePtr &plan,
        ReferenceQueryRunner *referenceQueryRunner);
} // namespace kumo::pollux::exec::test
